from scipy import *
from scipy import constants
from scipy import integrate
from numpy import *
from matplotlib.pyplot import *
import pylab as m

#Constants
#from scipy import hbar
hbar = 1.0 #Natural Units
BohrRadius=constants.physical_constants["Bohr radius"]
a0=BohrRadius[0]
z=(0+1j)#complex numbers: i

#controlling array size - currently picked for a convenient graph size
n=50
d=0.05
#declare arrays zero to avoid problems
f=zeros((n,n))


#Harmonic Oscillator (in 2d, natural units)

for i in range (0,n):
    for j in range  (0,n):
        
        x=i
        y=j
        
        r=d*((x**2 + y**2)**0.5)
        
        u_0 = (pi**(-1.0/2.0))*exp((-1.0*(r**2))/2.0)
        
        f[i,j] = u_0#setting array to wavefunction


##        # for second part - no longer using, currently left to remind me how to use quad!
##
##        E=0.5
##
##        V=0.5*r**2
##
##        K=E-V
##
##        L = K-V
##
##        l= lambda i : L
##
##        S = integrate.quad(l, 0, 100)
##
##        f_2[i,j] = S[0]
##
##
##s=zeros((n,n), 'complex')
##s = angle(f)    #gives argurement


#to produce graph

pcolor(f)
show()


